
# -*- coding: utf-8 -*-
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved

from fvcore.common.file_io import PathManager
import os
import numpy as np
import logging
import json

from detectron2.structures import BoxMode
from detectron2.data import DatasetCatalog, MetadataCatalog


__all__ = ["register_crowd_human"]


# fmt: off
CLASS_NAMES = [
    "hbox"
]
# fmt: on

def load_human_instances(dirname: str, split: str):
    """
    Load CrowdHuman detection annotations to Detectron2 format.

    Args:
        dirname: Contain "annotation_split.odgt"
        split (str): one of "train", "val"
    """
    logger = logging.getLogger(__name__)
    logger.info("Preprocessing CrowdHuman annotations ...")
    with open(os.path.join(dirname, "size.json")) as f:
        size = json.load(f)
    with open(os.path.join(dirname, f"annotation_{split}.odgt")) as f:
        annos = f.readlines()
    
    dicts = []
    ignored_count = 0
    for ann in map(json.loads, annos):
        jpeg_file = os.path.join(dirname, "Images", ann["ID"]+".jpg")
        r = {
            "file_name": jpeg_file,
            "image_id": ann["ID"],
            "height": size[ann["ID"]][0],
            "width": size[ann["ID"]][1],
        }
        instances = []
        for obj in ann['gtboxes']:
            boxes_diff = obj['tag'] != 'person'
            if not boxes_diff:
                if 'extra' in obj.keys():
                    extra = obj['extra']
                    if 'ignore' in extra.keys() and extra['ignore']:
                        boxes_diff = True
                    if 'unsure' in extra.keys() and extra['unsure']:
                        boxes_diff = True
            for cls in CLASS_NAMES:
                difficult = boxes_diff
                if 'hbox' == cls  and not difficult:
                    head_attr = obj['head_attr']
                    if 'ignore' in head_attr.keys() and head_attr['ignore'] or 'unsure' in head_attr.keys() and head_attr['unsure']:
                        difficult = True
                
                if difficult:
                    ignored_count +=1
                    continue
                x, y, w, h = obj[cls]
                bbox = [x, y, x+w, y+h]
                instances.append(
                    {"category_id": CLASS_NAMES.index(cls), "bbox": bbox, "bbox_mode": BoxMode.XYXY_ABS}
                )
        r["annotations"] = instances
        # if len(instances) >= 340 and split == 'train':
        #     continue
        dicts.append(r)
    logger.info(f"Load done! {ignored_count} Boxes are ignored!")
    return dicts


def register_crowd_human(name, dirname, split):
    DatasetCatalog.register(name, lambda: load_human_instances(dirname, split))
    MetadataCatalog.get(name).set(
        thing_classes=CLASS_NAMES, dirname=dirname, split=split
    )
